#include <iostream>
#include "Lottery.h"

#define TELEPORT_TRIES 1000

double IndexTree::Ask(int From,int To)
{
	double res = 0.0;
    int q,w;
	for(q=To+1;q>From;q-=RightBit(q))res+=Store[q-1];
	for (w=From;w>q;w-=RightBit(w))res-=Store[w-1];
	return res;
}

int LotteryRandomWalk::nodeToExplore() {
    if(unexploredCount <= 0) {
        int nextNode = network->getTeleportNode();
        //This could loop infinitely if we've explored everynode, unless we have cost limit
        for(int i = 0; visitedNode.insert(make_pair(nextNode, -1)).second == false && i < TELEPORT_TRIES; i++)
            nextNode = network->getTeleportNode();
        return nextNode;
    }
    double totalDeg = degreeCount.Ask(0, unexploredCount - 1);
    double val = double(rand30()) * totalDeg / double(RAND_MAX30);
    
    //TODO: discount edges from discovered node from ther degree in degreeCount
    //bsearch for r in index tree
    int l = 0, r = unexploredCount - 1;
    while(l < r) {
        int mid = (l + r) >> 1;
        if (degreeCount.Ask(0,mid) < val)
            l = mid + 1;
        else r = mid;
    }
    
    int retVal = nodeIDs[r];
    double tmp = -degreeCount.Ask(r,r);
    if(unexploredCount > 1) tmp += degreeCount.Ask(unexploredCount - 1, unexploredCount - 1);
    nodeIDs[r] = nodeIDs[--unexploredCount];
    degreeCount.Add(r, tmp);
    visitedNode[nodeIDs[r]] = r;
    visitedNode[retVal] = -1;
    return retVal;
}

void LotteryRandomWalk::updateDiscoveredNode(int discoveredNode) {
    if(visitedNode.find(discoveredNode) == visitedNode.end()) {
        TIntPr p = network->getDegree(discoveredNode);
        degreeCount.Add(unexploredCount, p.Val1 + p.Val2);
        if(unexploredCount == nodeIDs.size())
            nodeIDs.push_back(discoveredNode);
        else nodeIDs[unexploredCount] = discoveredNode;
        visitedNode[discoveredNode] = ++unexploredCount;
    }
}

// For testing purpose
#ifdef LOTTERY_TEST

int main(int argc, char* argv[]) {
    const int MAX_COST = 10000;
	PNGraph Gtmp = TSnap::GenForestFire(100, 0.35, 0.32);
	API *a = new API(TSnap::ConvertGraph<PNEGraph>(Gtmp));
	LotteryRandomWalk lott(a);
    //API net(1000);
    //LotteryRandomWalk lott(net);
    lott.setCostLimit(MAX_COST);
    while(lott.getTotalCost() < MAX_COST) {
        lott.nextStep();
        TIntPr p = lott.evaluate();
        cout << "Lottery cost used: " << lott.getTotalCost() << "  node recall: " << p.Val1  << endl;
    }
    return 0;
}

#endif
